!unzip -oq /home/aistudio/data/data103653/NEU-DET-VOC.zip -d /home/aistudio/
!pip install paddlex
import numpy as np
import matplotlib.pyplot as plt

def draw_bar(labels,quants):
    # X轴位置
    x = np.arange(len(labels))
    # 柱图大小
    width = 0.6
    # 创建图形
    fig, ax = plt.subplots()

    ax.bar(x+width, quants, width,color='blue')

    # X轴坐标显示，x + width*2 标识X轴刻度所在位置
    ax.set_xticks(x + width)
    ax.set_xticklabels(labels, rotation=45)
    # 显示右上角图例
    ax.legend()
    
    # 自动调整子图参数以提供指定的填充。多数情况下没看出来区别
    fig.tight_layout()
    
    plt.show()

from glob import glob
import pandas as pd
import numpy as np
import os
import cv2
from PIL import Image
from matplotlib import pyplot as plt
from tqdm import tqdm

#%%
TRAIN_DATASET_PATH = '/home/aistudio/NEU-DET-VOC/JPEGImages/'
# 训练集探索
labels = ['crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale', 'scratches']

# 每个类别的数量
dir_lst = os.listdir(TRAIN_DATASET_PATH)
quants  = []

for label in labels:
    path = os.path.join(TRAIN_DATASET_PATH,label)
    print(path)
    quant = len(glob(path+'*'))
    quants.append(quant)

draw_bar(labels, quants)

from PIL import Image, ImageEnhance, ImageDraw, ImageFont
import numpy as np
import matplotlib.pyplot as plt


plt.subplot(231)
plt.imshow(np.array(Image.open('/home/aistudio/NEU-DET-VOC/JPEGImages/crazing_1.jpg')))
plt.subplot(232)
plt.imshow(np.array(Image.open('/home/aistudio/NEU-DET-VOC/JPEGImages/inclusion_1.jpg')))
plt.subplot(233)
plt.imshow(np.array(Image.open('/home/aistudio/NEU-DET-VOC/JPEGImages/patches_1.jpg')))
plt.subplot(234)
plt.imshow(np.array(Image.open('/home/aistudio/NEU-DET-VOC/JPEGImages/pitted_surface_1.jpg')))
plt.subplot(235)
plt.imshow(np.array(Image.open('/home/aistudio/NEU-DET-VOC/JPEGImages/rolled-in_scale_1.jpg')))
plt.subplot(236)
plt.imshow(np.array(Image.open('/home/aistudio/NEU-DET-VOC/JPEGImages/scratches_1.jpg')))


!paddlex --split_dataset --format VOC --dataset_dir /home/aistudio/NEU-DET-VOC/ --val_value 0.15  --test_value 0.05

import paddlex as pdx
from paddlex.det import transforms
import matplotlib
matplotlib.use('Agg') 
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'

train_transforms = transforms.Compose([
        transforms.Normalize(),
        transforms.ResizeByShort(300, 500),
        transforms.Padding(32)
])

eval_transforms = transforms.Compose([
    transforms.Normalize(),
        transforms.ResizeByShort(300, 500),
        transforms.Padding(32)
])

train_dataset = pdx.datasets.VOCDetection(
                        data_dir='/home/aistudio/NEU-DET-VOC/', 
                        file_list='/home/aistudio/NEU-DET-VOC/train_list.txt',
                        label_list='/home/aistudio/NEU-DET-VOC/labels.txt',
                        transforms=train_transforms)
eval_dataset = pdx.datasets.VOCDetection(
                        data_dir='/home/aistudio/NEU-DET-VOC/',
                        file_list='/home/aistudio/NEU-DET-VOC/val_list.txt',
                        label_list='/home/aistudio/NEU-DET-VOC/labels.txt',
                        transforms=eval_transforms)

num_classes = len(train_dataset.labels)
model = pdx.det.PPYOLO(num_classes)
model.get_model_info()

model.train(
    num_epochs=40,
    train_dataset=train_dataset,
    train_batch_size=8,
    eval_dataset=eval_dataset,
    learning_rate=0.000125,
    lr_decay_epochs=[9, 14],
    save_interval_epochs=1,
    save_dir='output/PPYOLO_1',
    early_stop = False,
    )

model = pdx.load_model('/home/aistudio/output/PPYOLO_1/best_model/')  
result = model.evaluate(eval_dataset)
print(result)

test_dataset = pdx.datasets.VOCDetection(
                        data_dir='/home/aistudio/NEU-DET-VOC/',
                        file_list='/home/aistudio/NEU-DET-VOC/test_list.txt',
                        label_list='/home/aistudio/NEU-DET-VOC/labels.txt',
                        transforms=eval_transforms)
test_file_path = []

#读取test_list.txt
with open('/home/aistudio/NEU-DET-VOC/test_list.txt',"r") as tf:
    lines = tf.readlines()
    for line in lines:
        file_name = line.split(' ')[0]
        test_file_path.append('/home/aistudio/NEU-DET-VOC/{0}'.format(file_name))

model = pdx.load_model('/home/aistudio/output/PPYOLO_1/best_model/') 
result_list = model.batch_predict(test_file_path)
indexs = [2, 15, 38, 88]
for idx in indexs:
    pdx.det.visualize(test_file_path[idx], result_list[idx], threshold=0.1, save_dir='./')


model = pdx.load_model('/home/aistudio/output/PPYOLO_1/best_model/')  
image = '/home/aistudio/NEU-DET-VOC/JPEGImages/inclusion_273.jpg'
result = model.predict(image)
pdx.det.visualize(image, result, threshold=0.1, save_dir='./')

